import torch

from .distributed import Group
from .sum import Metric, SumMetric


class MeanMetric(Metric[torch.Tensor | None, torch.Tensor | tuple[torch.Tensor, torch.Tensor]]):
    def __init__(self, *, group: Group | None = None):
        self._value = SumMetric(group=group)
        self._weight = SumMetric(group=group)

    def compute(self) -> torch.Tensor | None:
        value = self._value.compute()
        weight = self._weight.compute()
        if value is None or weight is None:
            return None
        return value / weight

    def update(self, patch: torch.Tensor | tuple[torch.Tensor, torch.Tensor]):
        if isinstance(patch, tuple):
            value, weight = patch
        else:
            value, weight = patch, 1
        self._value.update(value)
        self._weight.update(weight)

    def reset(self):
        self._value.reset()
        self._weight.reset()
